import asyncio
import os
import signal

from fastapi import FastAPI, BackgroundTasks
from fastapi.responses import PlainTextResponse, Response
import uvicorn
import time
from multiprocessing import Process
from threading import Thread
import subprocess

app = FastAPI()

pids = {}


@app.get('/info', response_class=PlainTextResponse)
async def info():
    print('info')
    return 'info'


@app.get('/hello')
async def hello():
    print('hello')
    p = Process(target=stop, args=(20,))
    p.start()
    return Response(content='hello')


def stop(interval: int):
    for i in range(interval):
        time.sleep(1)
        print(Process.pid, "-->", str(i))


@app.post('/execute')
async def execute(task_name, interval):
    command = "for i in {1.." + interval + "}; do echo 'hello'; sleep 1; done"
    print(command)
    p = Thread(target=exec_command, args=(command, task_name))
    p.start()
    return 'execute success'


@app.post("/stop")
async def stopProcess(task_name):
    print(pids)
    pid = pids[task_name]
    try:
        os.kill(pid, signal.SIGTERM)  # 发送终止信号
        print(f"Process with PID {pid} has been terminated.")
    except ProcessLookupError:
        print(f"No process found with PID {pid}.")
    except PermissionError:
        print(f"Permission denied when trying to terminate process with PID {pid}.")
    except Exception as e:
        print(f"An error occurred: {e}")
    del pids[task_name]


def exec_command(command, task_name):
    for pid, line in run_command(command):
        print(pid, ': ', line)
        pids[task_name] = pid
        print(pids)


def run_command(command):
    process = subprocess.Popen(
        command,
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=1,
        universal_newlines=True
    )
    pid = process.pid

    with process.stdout:
        for line in iter(process.stdout.readline, ''):
            yield pid, line

    process.wait()
    return pid, process.returncode


@app.get("/async")
async def async_opt():
    await asyncio.sleep(20)
    return "async"


@app.get("/background_task")
async def back_task(background_tasks: BackgroundTasks):
    command = "for i in {1..20}; do echo 'hello'; sleep 1; done"
    # background_tasks.add_task(run_long_task, command)
    await run_long_task(command)
    return "background_task"


async def run_long_task(command):
    result = await execute_shell_command(command)
    print("execute done:")
    print(result)


async def execute_shell_command(command):
    print("execute_shell_command", command)
    process = await asyncio.create_subprocess_shell(
        command,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )
    # 实时读取输出
    while True:
        line = await process.stdout.readline()
        if not line:
            break
        print(line.decode().strip())
    await process.wait()

    return process.returncode


@app.post("/execute_async")
async def execute_async():
    """
    异步命令
    :return:
    """
    command = "for i in {1..20}; do echo 'hello'; sleep 1; done"
    await execute_train(command)


async def execute_train(command):
    def handle_log(pid, line):
        """
        回调处理函数
        """
        print(pid, ": ", line)

    await async_run_command(command, handle_log)


async def async_run_command(command, callback):
    """
    异步执行linux命令，并获取输出
    """
    # 使用 subprocess.Popen 执行命令
    process = await asyncio.create_subprocess_shell(
        command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT
    )
    pid = process.pid
    # 实时读取输出
    async for line in process.stdout:
        callback(pid, line.decode().strip())

    # 等待命令执行完成
    await process.wait()
    # 返回命令的退出状态
    return pid, process.returncode


if __name__ == '__main__':
    uvicorn.run(app, host='0.0.0.0', port=28000)
